import cv2
import numpy as np
import mediapipe as mp

# ---------------- Mediapipe 初始化 ----------------
mp_face_mesh = mp.solutions.face_mesh

def get_landmarks(image):
    """获取人脸关键点"""
    h, w, _ = image.shape
    with mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1,
                               refine_landmarks=True, min_detection_confidence=0.5) as face_mesh:
        results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        if not results.multi_face_landmarks:
            return None
        lm = results.multi_face_landmarks[0].landmark
        points = np.array([[int(p.x * w), int(p.y * h)] for p in lm], dtype=np.int32)
        return points

def get_delaunay_triangles(rect, points):
    """返回 Delaunay 三角形索引列表"""
    subdiv = cv2.Subdiv2D(rect)
    for p in points:
        subdiv.insert((float(p[0]), float(p[1])))  # 必须 float
    triangle_list = subdiv.getTriangleList()
    delaunay_tri = []
    for t in triangle_list:
        pts = [(t[0], t[1]), (t[2], t[3]), (t[4], t[5])]
        idx = []
        for pt in pts:
            for i, p in enumerate(points):
                if abs(pt[0] - p[0]) < 1 and abs(pt[1] - p[1]) < 1:
                    idx.append(i)
        if len(idx) == 3:
            delaunay_tri.append(tuple(idx))
    return delaunay_tri

def warp_triangle(img1, img2, t1, t2):
    """用仿射把三角形 t1 从 img1 warp 到 t2 在 img2"""
    r1 = cv2.boundingRect(np.float32([t1]))
    r2 = cv2.boundingRect(np.float32([t2]))

    t1_rect = []
    t2_rect = []
    for i in range(3):
        t1_rect.append(((t1[i][0] - r1[0]), (t1[i][1] - r1[1])))
        t2_rect.append(((t2[i][0] - r2[0]), (t2[i][1] - r2[1])))

    img1_rect = img1[r1[1]:r1[1]+r1[3], r1[0]:r1[0]+r1[2]]
    mask = np.zeros((r2[3], r2[2], 3), dtype=np.float32)
    cv2.fillConvexPoly(mask, np.int32(t2_rect), (1.0, 1.0, 1.0))

    M = cv2.getAffineTransform(np.float32(t1_rect), np.float32(t2_rect))
    warped = cv2.warpAffine(img1_rect, M, (r2[2], r2[3]), None,
                            flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)

    warped = warped * mask
    img2[r2[1]:r2[1]+r2[3], r2[0]:r2[0]+r2[2]] = \
        img2[r2[1]:r2[1]+r2[3], r2[0]:r2[0]+r2[2]] * (1 - mask) + warped

def face_swap(src_img, dst_img):
    """换脸主函数"""
    src_points = get_landmarks(src_img)
    dst_points = get_landmarks(dst_img)
    if src_points is None or dst_points is None:
        return dst_img

    h, w, _ = dst_img.shape
    rect = (0, 0, w, h)
    dt = get_delaunay_triangles(rect, dst_points)

    output = dst_img.copy()
    for tri in dt:
        t1 = [src_points[tri[0]], src_points[tri[1]], src_points[tri[2]]]
        t2 = [dst_points[tri[0]], dst_points[tri[1]], dst_points[tri[2]]]
        warp_triangle(src_img, output, t1, t2)

    # mask + 羽化
    hull = cv2.convexHull(dst_points)
    mask = np.zeros((h, w), dtype=np.uint8)
    cv2.fillConvexPoly(mask, hull, 255)
    mask = cv2.GaussianBlur(mask, (25, 25), 10)
    center = (w//2, h//2)

    # 自然融合
    output = cv2.seamlessClone(output, dst_img, mask, center, cv2.NORMAL_CLONE)
    return output

# ---------------- 实时摄像头换脸 ----------------
def main():
    cap = cv2.VideoCapture(0)
    src = cv2.imread("face_target.png")  # 源脸
    if src is None:
        print("请确保 face_source.jpg 存在！")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        result = face_swap(src, frame)

        cv2.imshow("Live Face Swap", result)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
